/***
 * Copyright (c) 2021-2031 murenchao
 * fig is licensed under Mulan PubL v2.
 * You can use this software according to the terms and conditions of the Mulan PubL v2.
 * You may obtain a copy of Mulan PubL v2 at:
 *       http://license.coscl.org.cn/MulanPubL-2.0
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 * See the Mulan PubL v2 for more details.
 */
package cool.taomu.software.fig.crypto.utils

import cool.taomu.software.fig.crypto.Base64
import java.io.ByteArrayInputStream
import java.io.ByteArrayOutputStream
import java.io.File
import java.nio.charset.Charset
import java.security.SecureRandom
import java.security.Security
import org.apache.commons.io.FileUtils
import org.bouncycastle.jce.provider.BouncyCastleProvider
import org.bouncycastle.openpgp.PGPCompressedData
import org.bouncycastle.openpgp.PGPCompressedDataGenerator
import org.bouncycastle.openpgp.PGPEncryptedData
import org.bouncycastle.openpgp.PGPEncryptedDataGenerator
import org.bouncycastle.openpgp.PGPEncryptedDataList
import org.bouncycastle.openpgp.PGPLiteralData
import org.bouncycastle.openpgp.PGPOnePassSignatureList
import org.bouncycastle.openpgp.PGPPublicKeyEncryptedData
import org.bouncycastle.openpgp.PGPPublicKeyRingCollection
import org.bouncycastle.openpgp.PGPSecretKeyRingCollection
import org.bouncycastle.openpgp.PGPSignatureList
import org.bouncycastle.openpgp.PGPUtil
import org.bouncycastle.openpgp.jcajce.JcaPGPObjectFactory
import org.bouncycastle.openpgp.operator.jcajce.JcaKeyFingerprintCalculator
import org.bouncycastle.openpgp.operator.jcajce.JcePBESecretKeyDecryptorBuilder
import org.bouncycastle.openpgp.operator.jcajce.JcePGPDataEncryptorBuilder
import org.bouncycastle.openpgp.operator.jcajce.JcePublicKeyDataDecryptorFactoryBuilder
import org.bouncycastle.openpgp.operator.jcajce.JcePublicKeyKeyEncryptionMethodGenerator
import org.bouncycastle.util.io.Streams
import org.slf4j.LoggerFactory

class PgpUtils {
    val static LOG = LoggerFactory.getLogger(PgpUtils);

    def static void main(String[] args) {
        Security.addProvider(new BouncyCastleProvider());
        var String pk = FileUtils.readFileToString(new File("src/main/resources/publicKey1.txt"),
            Charset.forName("UTF-8"));
        var data = PgpUtils.encrypt("Hello World!!!! ".bytes, pk, true);
        println(new String(data, "UTF-8"));
        var String prk = FileUtils.readFileToString(new File("src/main/resources/privateKey1.txt"),
            Charset.forName("UTF-8"));
        var d = PgpUtils.decrypt(data, prk, "123456789");
        println("result : " + new String(d, "UTF-8"));
    }

    def static decrypt(byte[] data, String privateKey, String password) {

        var bais = new ByteArrayInputStream(data);
        var input = PGPUtil.getDecoderStream(bais);
        var jpof = new JcaPGPObjectFactory(input);
        var o = jpof.nextObject;
        var PGPEncryptedDataList enc;
        if (o instanceof PGPEncryptedDataList) {
            enc = o as PGPEncryptedDataList;
        } else {
            enc = jpof.nextObject as PGPEncryptedDataList;
        }
        // 获取privatekey
        val pgpSec = new PGPSecretKeyRingCollection(new Base64(privateKey.bytes).decode(),
            new JcaKeyFingerprintCalculator());
        var edo = enc.encryptedDataObjects
        var pbe = edo.findFirst [ k |
            pgpSec.getSecretKey((k as PGPPublicKeyEncryptedData).keyID) !== null
        ] as PGPPublicKeyEncryptedData;
        var secKey = pgpSec.getSecretKey(pbe.keyID);
        var sKey = secKey.extractPrivateKey(
            new JcePBESecretKeyDecryptorBuilder().setProvider("BC").build(password.toCharArray))
        var clear = pbe.getDataStream(new JcePublicKeyDataDecryptorFactoryBuilder().setProvider("BC").build(sKey));
        var plain = new JcaPGPObjectFactory(clear);
        var message = plain.nextObject
        try(var out = new ByteArrayOutputStream()) {
            if (message instanceof PGPCompressedData) {
                var cdata = message as PGPCompressedData;
                var fact = new JcaPGPObjectFactory(cdata.dataStream);
                message = fact.nextObject;
            }
            if (message instanceof PGPLiteralData) {
                var id = message as PGPLiteralData;
                Streams.pipeAll(id.inputStream, out);
            } else if (message instanceof PGPOnePassSignatureList) {
                LOG.info("encrypted message contains a signed message - not literal data.")
            } else if (message instanceof PGPSignatureList) {
                LOG.info("encrypted message contains a signed message - not literal data.");
            }
            if (pbe.isIntegrityProtected) {
                if (!pbe.verify) {
                    LOG.info("message failed integrity check");
                } else {
                    LOG.info("message integrity check passed");
                }
            } else {
                LOG.info("no message integrity check");
            }
            return out.toByteArray
        }
    }

    def static encrypt(byte[] data, String publicKey, boolean withIntegrityCheck) {
        // 获取publiekey
        var bais = new ByteArrayInputStream(new Base64(publicKey.bytes).decode);
        var input = PGPUtil.getDecoderStream(bais);
        var pkrc = new PGPPublicKeyRingCollection(input, new JcaKeyFingerprintCalculator);
        var key = pkrc.keyRings.findFirst [ kr |
            var pk = kr.publicKeys.findFirst[pk|pk.isEncryptionKey]
            pk.isEncryptionKey;
        ]
        var encKey = key.publicKey;

        var file = CreateTempFile.create(data, "UTF-8");
        try(var baos = new ByteArrayOutputStream()) {
            var cdg = new PGPCompressedDataGenerator(PGPCompressedData.ZIP);
            PGPUtil.writeFileToLiteralData(cdg.open(baos), 'b', file);
            cdg.close;
            var bytes = baos.toByteArray;
            var encGen = new PGPEncryptedDataGenerator(
                new JcePGPDataEncryptorBuilder(PGPEncryptedData.CAST5).setWithIntegrityPacket(withIntegrityCheck).
                    setSecureRandom(new SecureRandom()).setProvider("BC"));
            encGen.addMethod(new JcePublicKeyKeyEncryptionMethodGenerator(encKey).provider = "BC");
            try(var out = new ByteArrayOutputStream()) {
                var cout = encGen.open(out, bytes.length);
                cout.write(bytes);
                cout.close();
                return out.toByteArray;
            }

        } finally {
            file.delete
        }
    }
}
